import trimesh
import numpy as np
import json
from scipy.spatial.transform import Rotation
# import torch
from pathlib import Path
import sys
from save_file import save_info
from scalePcd import scalePcd   
import os
import argparse
from pathlib import Path

def glb_to_pcd(glb_path, sample_num=None):
    mesh = trimesh.load(glb_path)
    if isinstance(mesh, trimesh.Scene):
        mesh = list(mesh.geometry.values())[0]
    vertices = np.array(mesh.vertices)
    colors = np.ones((len(vertices), 3)) * 0.5
    if sample_num and len(vertices) > sample_num:
        indices = np.random.choice(len(vertices), sample_num, replace=False)
        vertices = vertices[indices]
        colors = colors[indices]
    return np.concatenate([vertices, colors], axis=1)

# def torch_to_np(tensor):
#     if tensor is None:
#         return None
#     if isinstance(tensor, torch.Tensor):
#         if tensor.is_cuda:
#             tensor = tensor.cpu().detach()
#         return tensor.numpy()
#     elif not isinstance(tensor, np.ndarray):
#         tensor = np.array(tensor)
#     return tensor

def PointArrow(mat, color, pointnum=100,show_xy_arrow=False):
    mat = np.array(mat)
    if mat.shape == (6,) or (len(mat.shape) == 1 and mat.shape[0] == 6):
        pos = mat[:3] 
        roll, pitch, yaw = mat[3:]
        rotation = Rotation.from_euler('xyz', [roll, pitch, yaw])
        x_direction = rotation.apply([1, 0, 0])
        y_direction = rotation.apply([0, 1, 0])
        z_direction = rotation.apply([0, 0, 1])
    elif mat.shape == (7,) or (len(mat.shape) == 1 and mat.shape[0] == 7):
        pos = mat[:3] 
        quat = mat[3:] 
        rotation = Rotation.from_quat([quat[1], quat[2], quat[3], quat[0]])
        x_direction = rotation.apply([1, 0, 0])
        y_direction = rotation.apply([0, 1, 0])
        z_direction = rotation.apply([0, 0, 1])
    elif mat.shape == (4, 4):
        pos = mat[:3, 3]
        x_direction = mat[:3, 0]
        y_direction = mat[:3, 1]
        z_direction = mat[:3, 2]
    else:
        raise ValueError(f"PointArrow expects 4x4 matrix or 7-element pose list, got shape {mat.shape}")
    
    def create_arrow(direction, arrow_length, arrow_color):
        line_points = np.linspace(0, arrow_length, pointnum//2)
        line_vertices = np.array([pos + t * direction for t in line_points]).reshape(-1, 3)
        ball_center = pos + arrow_length * direction
        ball_radius = 0.005
        ball_points = ball_center + ball_radius * np.random.randn(pointnum//2, 3)
        vertices = np.concatenate([line_vertices, ball_points])
        line_colors = np.ones((len(line_vertices), 3)) * arrow_color
        ball_colors = np.ones((len(ball_points), 3))
        colors = np.concatenate([line_colors, ball_colors], axis=0)
        return np.concatenate([vertices, colors], axis=1)
    
    # z_arrow = create_arrow(z_direction, 0.12, color)
    # return np.concatenate([z_arrow], axis=0)
    if show_xy_arrow:
        x_arrow = create_arrow(x_direction, 0.04, [0, 0, 1])  # Blue, shorter
        y_arrow = create_arrow(y_direction, 0.04, [0, 1, 0])  # Green, shorter
    else:
        x_arrow = np.empty((0, 6))
        y_arrow = np.empty((0, 6))
    z_arrow = create_arrow(z_direction, 0.08, color)
    return np.concatenate([x_arrow, y_arrow, z_arrow], axis=0)

def create_annotation_pcd(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    annotations = []
    for pose in data.get('target_pose', []):
        annotations.append(PointArrow(pose, [1, 0, 0], 20))
    for pose in data.get('contact_points_pose', []):
        annotations.append(PointArrow(pose, [0, 1, 0], 20))
    for pose in data.get('functional_matrix', []):
        annotations.append(PointArrow(pose, [0, 0, 1], 20))
    return np.concatenate(annotations, axis=0) if annotations else np.empty((0, 6))

def catpcd(pcd1, pcd2, color2=[0, 1, 0], formatstr='xc'):
    pcd2 = np.array(pcd2)
    assert len(pcd2.shape) == 2, f"pcd2 must be 2D array"
    assert pcd2.shape[1] == len(formatstr) * 3, f"pcd2 must be {len(formatstr) * 3} dim for format {formatstr}"
    pcd1 = scalePcd(pcd1)
    pcd2 = scalePcd(pcd2)
    if pcd2.shape[0] == 0:
        return pcd1
    if 'r' in formatstr:
        x_idx, r_idx = formatstr.index('x') * 3, formatstr.index('r') * 3
        c_idx = formatstr.index('c') * 3 if 'c' in formatstr else None
        pcd2 = np.concatenate([PointArrow(np.concatenate([row[x_idx:x_idx+3], row[r_idx:r_idx+3]]), row[c_idx:c_idx+3] if c_idx is not None else color2) for row in pcd2])
    if formatstr == 'x':
        pcd2 = np.concatenate([pcd2, np.ones((pcd2.shape[0], 3)) * color2], axis=1)
    return np.concatenate([pcd1, pcd2], axis=0)

def contact12_gripper36_to_pcd(contact12, gripper36, origin_pcd=None):
    # contact12 = torch_to_np(contact12)
    # gripper36 = torch_to_np(gripper36)
    # origin_pcd = torch_to_np(origin_pcd)
    pcd = np.zeros((0, 6))
    contact_points = contact12.reshape(2, 6)
    pcd = catpcd(pcd, contact_points, color2=[0, 1, 0], formatstr='xr')
    pcd = catpcd(pcd, gripper36[0:6].reshape(1, 6), color2=[1, 0, 0], formatstr='xr')
    pcd = catpcd(pcd, gripper36[6:18].reshape(4, 3), color2=[0, 1, 0], formatstr='x')
    pcd = catpcd(pcd, gripper36[18:24].reshape(1, 6), color2=[0, 0.5, 1], formatstr='xr')
    pcd = catpcd(pcd, gripper36[24:36].reshape(4, 3), color2=[1, 0.5, 0], formatstr='x')
    if origin_pcd is not None:
        if origin_pcd.shape[1] == 3:
            pcd = catpcd(pcd, origin_pcd, color2=[1, 1, 1], formatstr='x')
        elif origin_pcd.shape[1] == 6:
            pcd = catpcd(pcd, origin_pcd, formatstr='xc')
        else:
            print(f"contact12_gripper36_to_pcd optional origin_pcd is irregulary shaped {origin_pcd.shape}, can not cat with existing pcd: {pcd.shape}")
    return pcd

def keypose48_to_pcd(keypose48):
    return contact12_gripper36_to_pcd(keypose48[:12], keypose48[12:])

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("object_name", type=str)
    args = parser.parse_args()
    basepath = f"{Path(__file__).parent.parent.parent}/assets/objects/{args.object_name}"
    pcd = glb_to_pcd(os.path.join(basepath, "visual/base0.glb"), 3000)
    annotation_pcd = create_annotation_pcd(os.path.join(basepath, "model_data0.json"))
    pcd = catpcd(pcd, annotation_pcd)
    save_info("pcd", os.path.join(basepath, "annotated_n_model_data0"), pcd)
